﻿#pragma once
#include<iostream>
#include<assert.h>
using namespace std;

template<class K, class V>
struct AVLTreeNode
{
	pair<K, V> _kv;
	AVLTreeNode<K, V>* _left;
	AVLTreeNode<K, V>* _right;
	AVLTreeNode<K, V>* _parent;
	int _bf;//平衡因子
	AVLTreeNode(const pair<K, V>& kv)
		:_kv(kv)
		,_left(nullptr)
		,_right(nullptr)
		,_parent(nullptr)
		,_bf(0)
	{}
};

template<class K, class V>
class AVLTree
{

	typedef AVLTreeNode<K, V> Node;
public:
	bool Insert(const pair<K, V> kv)
	{
		if (_root == nullptr)
		{
			_root = new Node(kv);
			return true;
		}
			Node* cur = _root;
			Node* parent = nullptr;
			while (cur)
			{
				if (cur->_kv.first < kv.first)
				{
					parent = cur;
					cur = cur->_right;
				}
				else if (cur->_kv.first > kv.first)
				{
					parent = cur;
					cur = cur->_left;
				}
				else
				{
					return false;
				}
			}
			cur = new Node(kv);
			if (parent->_kv.first < cur->_kv.first)
			{
				parent->_right = cur;
			}
			else
			{
				parent->_left = cur;
			}
			cur->_parent = parent;

			//更新平衡因子
			while (parent)
			{
				if (parent->_left == cur)
				{
					parent->_bf--;
				}
				else
				{
					parent->_bf++;
				}
				if (parent->_bf == 0)
					break;
				else if (parent->_bf == 1 || parent->_bf == -1)
				{
					//继续向上更新平衡因子
					cur = parent;
					parent = parent->_parent;
				}
				else if (parent->_bf == 2 || parent->_bf == -2)
				{
					//旋转逻辑
					if (parent->_bf == -2 && cur->_bf == -1)
					{
						RotateR(parent);
					}
					else if (parent->_bf == 2 && cur->_bf == 1)
					{
						RotateL(parent);
					}
					else if (parent->_bf == -2 && cur->_bf == 1)
					{
						RotateLR(parent);
					}
					else
					{
						RotateRL(parent);
					}
					break;
				}
				else
				{
					assert(false);
				}
			}
			return true;
	}

		void RotateR(Node * parent)
		{
			Node* subL = parent->_left;
			Node* subLR = subL->_right;
			//除了修改孩子指针指向，还要修改父亲指针
			parent->_left = subLR;
			if (subLR)
				subLR->_parent = parent;

			Node* parentParent = parent->_parent;
			subL->_right = parent;
			parent->_parent = subL;

			// parent有可能是整棵树的根，也可能是局部的⼦树
		  // 如果是整棵树的根，要修改_root
		  // 如果是局部的指针要跟上⼀层链接
			if (parentParent == nullptr)
			{
				_root = subL;
				subL->_parent = nullptr;
			}
			else
			{
				if (parent == parentParent->_left)
				{
					parentParent->_left = subL;
				}
				else
				{
					parentParent->_right = subL;
				}
				subL->_parent = parentParent;
			}

			parent->_bf = subL->_bf = 0;
		}

		void RotateL(Node * parent)
		{
			Node* subR = parent->_right;
			Node* subRL = subR->_left;
			parent->_right = subRL;
			if (subRL)
				subRL->_parent = parent;
			Node* parentParent = parent->_parent;
			subR->_left = parent;
			parent->_parent = subR;
			if (parentParent == nullptr)
			{
				_root = subR;
				subR->_parent = nullptr;
			}
			else
			{
				if (parent == parentParent->_left)
				{
					parentParent->_left = subR;
				}
				else
				{
					parentParent->_right = subR;
				}
				subR->_parent = parentParent;
			}
			parent->_bf = subR->_bf = 0;
		}

		void RotateLR(Node * parent)
		{
			Node* subL = parent->_left;
			Node* subLR = subL->_right;
			int bf = subLR->_bf;
			RotateL(parent->_left);
			RotateR(parent);
		
			if (bf == -1)
			{
				subLR->_bf = subL->_bf = 0;
				parent->_bf = 1;
			}
			else if (bf == 1)
			{
				subLR->_bf = parent->_bf = 0;
				subL->_bf = -1;
			}
			else if (bf == 0)
			{
				subLR->_bf = subL->_bf = parent->_bf = 0;
			}
			else
			{
				assert(false);
			}
		}


		void RotateRL(Node * parent)
		{
			Node* subR = parent->_right;
			Node* subRL = subR->_left;
			int bf = subRL->_bf;
			RotateR(parent->_right);
			RotateL(parent);
		
			if (bf == -1)
			{

				subRL->_bf = parent->_bf = 0;
				subR->_bf = 1;
			}
			else if (bf == 1)
			{
				subRL->_bf = subR->_bf = 0;
				parent->_bf = -1;
			}
			else if (bf == 0)
			{
				subR->_bf = subRL->_bf = parent->_bf = 0;
			}
			else
			{
				assert(false);
			}
		}

		Node* Find(const K & key)
		{
			Node* cur = _root;
			while (cur)
			{
				if (cur->_kv.first < key)
					cur = cur->_right;
				else if (cur->_kv.first > key)
					cur = cur->_left;
				else
					return cur;
			}
			return nullptr;
		}

		int Height()
		{
			return _Height(_root);
		}

		int Size()
		{
			return _Size(_root);
		}

		void InOrder()
		{
			 _InOrder(_root);
		}

		bool IsBalanceTree()
		{
			return _IsBalanceTree(_root);
		}



	private:
		int _Height(Node * root)
		{
			if (root == nullptr)
				return 0;
			int leftHeight = _Height(root->_left);
			int rightHeight = _Height(root->_right);
			return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
		}

		void _InOrder(Node * root)
		{
			if (root == nullptr)
			{
				return;
			}

			_InOrder(root->_left);
			cout << root->_kv.first << ":" << root->_kv.second << endl;
			_InOrder(root->_right);
		}

		int _Size(Node * root)
		{
			if (root == nullptr)
				return 0;

			return _Size(root->_left) + _Size(root->_right) + 1;
		}

		bool _IsBalanceTree(Node * root)
		{
			// 空树也是AVL树
			if (root == nullptr)
				return true;

			int leftHeight = _Height(root->_left);
			int rightHeight = _Height(root->_right);
			int diff = rightHeight - leftHeight;

			if (abs(diff) >= 2)
			{
				cout << root->_kv.first << "高度差异常" << endl;
				return false;
			}

			if (root->_bf != diff)
			{
				cout << root->_kv.first << "平衡因子异常" << endl;
				return false;
			}

			return _IsBalanceTree(root->_left) && _IsBalanceTree(root->_right);
		}

	private:
		Node* _root = nullptr;
};
